from sklearn.datasets import make_moons
from matplotlib import pyplot
from matplotlib.colors import ListedColormap
import numpy as np
import os
from DANN2 import DANN
def main():
    np.random.seed(100)
    # np.random.seed(47)
    import shutil
    # shutil.rmtree('dann_scalar')  # shutil.rmtree()表示递归删除文件夹下的所有子文件夹和子文件
    # os.mkdir('dann_scalar')
    # Xs,Ys,Xt,Yt = make_trans_moons(theta=30,number=200)
    #200 40 93.5#300 30 96 #200 20 99.5#320 20
    Xs, Ys, Xt, Yt = make_trans_moons(theta=30, number=200)
    print(Xs.shape, Ys.shape, Xt.shape)
    dann = DANN(Xs=Xs, Ys=Ys, Xt=Xt, Yt=Yt, save=False)
    dann.fit()
    # dann = DANN(Xs=Xs, Ys=Ys, Xt=Xt, Yt=Yt, save=True)

    pyplot.xticks([])
    pyplot.yticks([])
    pyplot.title("DITA classification")
    draw_trans_data(Xs, Ys, Xt, dann.predict)
    pyplot.savefig('tmp.pdf')
    pyplot.show()

    pyplot.close()


def make_trans_moons(theta=30, number=100, noise=.05):
    from math import cos, sin, pi
    Xs, Ys = make_moons(number, noise=noise, random_state=1)
    Xt, Yt = make_moons(number, noise=noise, random_state=2)

    trans = -np.mean(Xs, axis=0)
    Xs = 2 * (Xs + trans)
    Xt = 2 * (Xt + trans)
    #旋转
    theta = - theta * pi / 180
    rotation = np.array([[cos(theta), sin(theta)], [-sin(theta), cos(theta)]])
    Xt = np.dot(Xt, rotation.T)
    return Xs, Ys, Xt, Yt


def draw_trans_data(Xs,Ys,Xt,predict=None):
    # cm_bright = ListedColormap(['#FF0000','#33FFFF','#00FF00'])
    # cm_bright = ListedColormap(['#FF0000','#00FF00'])
    cm_bright = ListedColormap(['#9999CC', '#FFFF99'])

    # cm_bright = ListedColormap(['#FFFC33','#33FFF3'])#FFFF66 33FFF3 #4FD5D6
    x_min, x_max = 1.1 * Xs[:, 0].min(), 1.1 * Xs[:, 0].max()
    y_min, y_max = 1.5 * Xs[:, 1].min(), 1.5 * Xs[:, 1].max()

    pyplot.xlim((x_min, x_max))
    pyplot.ylim((y_min, y_max))

    pyplot.tick_params(direction='in', labelleft=False)

    if predict is not None:
        h = .02
        x, y = np.meshgrid(np.arange(x_min, x_max, h),
                          np.arange(y_min, y_max, h))
        # Z1,Z2 = predict(np.c_[x.ravel(),y.ravel()])
        # Z1 = Z1.reshape(x.shape).numpy()
        # Z2 = Z2.reshape(x.shape).numpy()
        # Z = Z1+Z2
        Z = predict(np.c_[x.ravel(), y.ravel()])
        Z = Z.reshape(x.shape).numpy()
        pyplot.contourf(x, y, Z, cmap=cm_bright, alpha=.4)
        pyplot.contour(x, y, Z, colors='black', linewidths=2)

    if Xs is not None:
        for i in range(len(Ys)):
            if Ys[i] == 1:
                pyplot.annotate(".", Xs[i, :], color="#FFCC33", size=40 * 1.5)
            else:
                pyplot.annotate(".", Xs[i, :], color="#9966FF", size=40 * 1.5) # red

    if Xt is not None:
        select = [3,111,88,53,16,181,163,167,33,188]
        xt_select = Xt[select]
        yt_select = Xt[select]
        pyplot.scatter(Xt[:, 0], Xt[:, 1], c='#3399FF', s=40)
        pyplot.scatter(xt_select[:, 0], yt_select[:, 1], c='red', s=100, marker='^')

if __name__ == '__main__':
    main()

